import os
def get_qwen3_client():
    from util_keys import DASHSCOPE_KEY
    from openai import OpenAI 
    client = OpenAI(
        # 若没有配置环境变量，请用百炼API Key将下行替换为：api_key="sk-xxx",
        api_key=DASHSCOPE_KEY,
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    )
    return client

def build_question(template_json_str, search_txt):
    """
    构造模型问题提示词
    - 检索语句需自动填入，json内容仅为结构参考，非默认值。
    - 未明确内容置空，不猜测。count未指定为-1。
    """
    return (
        f"{template_json_str} 检索语句：{search_txt} ."
        "先把检索语句转换为英文描述,然后"
        "将检索语句转换为为上面的json模板格式并输出，不要加入markdown语法，只是按这个输出即可。"
        "注意：json模板的内容只是参考，不要作为默认值！"
        "注意：检索语句中没有提到的内容全部填充为空，不要臆想！"
        "提示：关于count的部分，如果指定数量则按照指定的写, 若数量没有指定的情况下，then"
        "如果衣服类型有袖子，则数量默认是2，如果是裤子等没有袖子的衣服类型则默认为-1，"
        "口袋数量若没有指定，则填写-1"
        "纽扣数量若没有指定，则填写-1"
        "特殊的有无问题情况："
        "如果输入 有口袋 ，则口袋数量填写1，没有口袋 or 无口袋，填写0"
        "如果输入 有袖子，则袖子数量填写2，没有袖子 or 无袖，填写0"
        "如果输入 有纽扣，则纽扣数量填写1，没有纽扣 or 无纽扣，填写0"
    )
def is_valid_value(v):
    """
    判定value是否有效（非None、非-1、非空字符串/仅空白、非空列表）
    """
    if (
        v is None
        or v == ""
        or v == []
        or (isinstance(v, str) and (v.strip() == "" or v.lower() == "none" or v == "-1"))
        or v == -1
    ):
        return False
    return True
def extract_non_empty_keys(obj, parent_keys=None):
    """
    递归提取所有 value 满足 is_valid_value 的 key链
    返回格式: [ ["a"], ["b","c"], ...]，每个为一条有效路径
    """
    parent_keys = parent_keys or []
    results = []

    if isinstance(obj, dict):
        for k, v in obj.items():
            path = parent_keys + [k]
            if isinstance(v, dict):
                results.extend(extract_non_empty_keys(v, path))
            elif isinstance(v, list):
                # 若list元素是dict，递归每一项，收集子key链
                list_results = []
                tmp_list = []
                for idx, item in enumerate(v):
                    if isinstance(item, dict):
                        list_results += extract_non_empty_keys(item, path + [str(idx)])
                    else:
                        if is_valid_value(item):
                            tmp_list.append(item)
                if list_results:
                    results.extend(list_results)
                # list整体有效，添加当前链
                if is_valid_value(tmp_list):
                    results.append(path)
            else:
                if is_valid_value(v):
                    results.append(path)
    else:
        if is_valid_value(obj):
            results.append(parent_keys)
    return results
def get_value_by_keys(d, keys):
    """
    根据key链依次从嵌套dict/list中取值
    keys: 例如 ['design','pockets','0','count']
    若找不到返回None
    """
    cur = d
    for k in keys:
        if isinstance(cur, dict):
            if k in cur:
                cur = cur[k]
            else:
                return None
        elif isinstance(cur, list):
            try:
                idx = int(k)
                if 0 <= idx < len(cur):
                    cur = cur[idx]
                else:
                    return None
            except Exception:
                return None
        else:
            return None
    return cur
def get_ans_dict_by_search(qwen_client, template_json_str, search_txt):
    question = build_question(template_json_str, search_txt)

    ans = get_ans_by_client(qwen_client, question)

    return ans
def get_non_empty_keys(ans_dict):
    non_empty_keys = []
    if ans_dict is not None:
        non_empty_keys = extract_non_empty_keys(ans_dict)
        # key_strs = ["->".join(keys) for keys in non_empty_keys]
        # print("有效内容的key链:", key_strs)
    return non_empty_keys
def get_matched_info_by_jsondict(json_dict:dict, non_empty_keys:list):
    value_strs = []
    for keys in non_empty_keys:
        value = get_value_by_keys(json_dict, keys)
        if isinstance(value, list):
            val_str = ",".join(str(v) for v in value)
        else:
            val_str = str(value) if value is not None else ""
        value_strs.append(val_str)
        # print(f"path: {'->'.join(keys)}  value: {val_str}")
    # joined_values = "&&".join(value_strs)
    return value_strs
def get_matched_info(target_json_dict:dict, 
                    ans_dict:dict, non_empty_keys:list):
    """
    主流程：
    1. 加载模板json
    2. 生成待检索问题字符串，调用qwen3接口获得AI回复json字符串
    3. 解析AI回复为dict，提取所有“有效内容”的key链（多层键路径列表）
    4. 分别从模板和AI回复，通过这些key链提取最终字符串特征用于后续检索
    """
    # —— 从模板取内容并拼接
    # print("根据有效key链，获得template中对应内容：")
    joined_values_list = get_matched_info_by_jsondict(target_json_dict,
                                                non_empty_keys)
    joined_values = '&&'.join(joined_values_list)
    # print("收集的所有value并拼接:", joined_values)

    search_joined_values_list = get_matched_info_by_jsondict(ans_dict,
                                                        non_empty_keys)
    search_joined_values = '&&'.join(search_joined_values_list)
    # print("从ans_dict收集的所有value并拼接:", search_joined_values)

    # 返回最终字符串
    return search_joined_values, joined_values

def get_ans_by_client(client, question):
    completion = client.chat.completions.create(
        model="qwen3-235b-a22b-instruct-2507",
        messages=[
            {"role": "system", "content": "You are a helpful assistant for generating json data."},
            {"role": "user", "content": question},
        ],
        stream=True
    )
    result_chunks = []
    for chunk in completion:
        content = chunk.choices[0].delta.content
        # print(content,end='',flush=True)
        if content is not None:
            result_chunks.append(content)
    return "".join(result_chunks)


def test():
    """
    主流程：
    1. 加载模板json
    2. 生成待检索问题字符串，调用qwen3接口获得AI回复json字符串
    3. 解析AI回复为dict，提取所有“有效内容”的key链（多层键路径列表）
    4. 分别从模板和AI回复，通过这些key链提取最终字符串特征用于后续检索
    """
    import json

    # ========== Step 1: 读取模板json字符串 ==========
    template_json_path = 'demo_qwenvl_clothing.json'
    with open(template_json_path, 'r', encoding='utf-8') as f:
        template_json_str = f.read()

    # ========== Step 2: 准备输入提问 & 获取AI返回的答案 ==========
    def build_question(search_txt):
        """
        构造模型问题提示词
        - 检索语句需自动填入，json内容仅为结构参考，非默认值。
        - 未明确内容置空，不猜测。count未指定为-1。
        """
        return (
            f"{template_json_str} 检索语句：{search_txt} ."
            "将检索语句转换为为上面的json模板格式并输出，不要加入markdown语法，只是按这个输出即可。"
            "注意：json模板的内容只是参考，不要作为默认值！"
            "注意：检索语句中没有提到的内容全部填充为空，不要臆想！关于count的部分，没有指定数量，则填写-1."
        )

    search_txt = '春秋两季黑色大衣，两个口袋'
    question = build_question(search_txt)

    qwen_client = get_qwen3_client()
    ans = get_ans_by_client(qwen_client, question)
    # print(ans)  # 调试用

    # ========== Step 3: 解析AI返回为dict ==========
    try:
        ans_dict = json.loads(ans)
    except Exception as e:
        print("Failed to parse ans as JSON:", e)
        ans_dict = None
    # print(ans_dict)  # 调试用

    # ========== Step 4: 提取所有有效value的key链 ==========
    def is_valid_value(v):
        """
        判定value是否有效（非None、非-1、非空字符串/仅空白、非空列表）
        """
        if v is None or v == "" or v == [] or v == -1:
            return False
        if isinstance(v, str) and not v.strip():
            return False
        return True

    def extract_non_empty_keys(obj, parent_keys=None):
        """
        递归提取所有 value 满足 is_valid_value 的 key链
        返回格式: [ ["a"], ["b","c"], ...]，每个为一条有效路径
        """
        parent_keys = parent_keys or []
        results = []

        if isinstance(obj, dict):
            for k, v in obj.items():
                path = parent_keys + [k]
                if isinstance(v, dict):
                    results.extend(extract_non_empty_keys(v, path))
                elif isinstance(v, list):
                    # 若list元素是dict，递归每一项，收集子key链
                    list_results = []
                    for idx, item in enumerate(v):
                        if isinstance(item, dict):
                            list_results += extract_non_empty_keys(item, path + [str(idx)])
                    if list_results:
                        results.extend(list_results)
                    # list整体有效，添加当前链
                    elif is_valid_value(v):
                        results.append(path)
                else:
                    if is_valid_value(v):
                        results.append(path)
        else:
            if is_valid_value(obj):
                results.append(parent_keys)
        return results

    non_empty_keys = []
    if ans_dict is not None:
        non_empty_keys = extract_non_empty_keys(ans_dict)
        key_strs = ["->".join(keys) for keys in non_empty_keys]
        print("有效内容的key链:", key_strs)

    # ========== Step 5: 从模板和AI结果中通过key链取内容，再拼接 ==========
    def get_value_by_keys(d, keys):
        """
        根据key链依次从嵌套dict/list中取值
        keys: 例如 ['design','pockets','0','count']
        若找不到返回None
        """
        cur = d
        for k in keys:
            if isinstance(cur, dict):
                if k in cur:
                    cur = cur[k]
                else:
                    return None
            elif isinstance(cur, list):
                try:
                    idx = int(k)
                    if 0 <= idx < len(cur):
                        cur = cur[idx]
                    else:
                        return None
                except Exception:
                    return None
            else:
                return None
        return cur

    # 解析模板json字符串为dict
    try:
        template_json = json.loads(template_json_str)
    except Exception as e:
        print("template_json_str 解析失败:", e)
        template_json = {}

    # —— 从模板取内容并拼接
    print("根据有效key链，获得template中对应内容：")
    value_strs = []
    for keys in non_empty_keys:
        value = get_value_by_keys(template_json, keys)
        if isinstance(value, list):
            val_str = ",".join(str(v) for v in value)
        else:
            val_str = str(value) if value is not None else ""
        value_strs.append(val_str)
        print(f"path: {'->'.join(keys)}  value: {val_str}")
    joined_values = "&&".join(value_strs)
    print("收集的所有value并拼接:", joined_values)

    # —— 从AI结果取内容并拼接
    search_value_strs = []
    for keys in non_empty_keys:
        value = get_value_by_keys(ans_dict, keys)
        if isinstance(value, list):
            val_str = ",".join(str(v) for v in value)
        else:
            val_str = str(value) if value is not None else ""
        search_value_strs.append(val_str)
        print(f"[ori] path: {'->'.join(keys)}  value: {val_str}")
    search_joined_values = "&&".join(search_value_strs)
    print("从ans_dict收集的所有value并拼接:", search_joined_values)

    # 返回最终字符串
    return search_joined_values, joined_values

def test_matched_info():
    # ========== Step 1: 读取模板json字符串 ==========
    template_json_path = 'demo_qwenvl_clothing.json'
    with open(template_json_path, 'r', encoding='utf-8') as f:
        template_json_str = f.read()

    # ========== Step 2: 构建question + client + template = ans ==========
    search_txt = '春秋两季黑色大衣，两个口袋'

    qwen_client = get_qwen3_client()
    
    ans = get_ans_dict_by_search(qwen_client, template_json_str, search_txt)
    
    # ========== Step 3: ans => ans_dict ==========
    import json
    try:
        ans_dict = json.loads(ans)
    except Exception as e:
        print("Failed to parse ans as JSON:", e)
        ans_dict = None

    # ========== Step 4: ans_dict => non keys ==========
    non_empty_keys = get_non_empty_keys(ans_dict)

    # ========== Step 5: 构建 target_json_dict ==========
    try:
        target_json_dict = json.loads(template_json_str)
    except Exception as e:
        print("template_json_str 解析失败:", e)
        target_json_dict = {}

    # ========== Step 6: 获取 target_json_dict's concat value + ans_dict's concat value ==========
    # —— 从模板取内容并拼接
    search_joined_values_list = get_matched_info_by_jsondict(ans_dict,
                                                        non_empty_keys)
    search_joined_values = '&&'.join(search_joined_values_list)                                        
    print("从ans_dict收集的所有value并拼接:", search_joined_values)

    # —— 从模板取内容并拼接
    joined_values_list = get_matched_info_by_jsondict(target_json_dict,
                                                non_empty_keys)
    joined_values = '&&'.join(joined_values_list)
    print("收集的所有value并拼接:", joined_values)

    # from clip import embdding
    # 计算 base_sim = empty_txt_emb                 COSSIM joined_values_emb
    # 计算 search_sim = search_joined_values_emb    COSSIM joined_values_emb
    # relative_sim = search_sim - base_sim
    # if relative_sim < 0 : drop
    # else : 构建 top k(k=10)的 堆，按照 relative_sim 建堆
    from util_for_clipvitb32 import clip_vit_b32, get_embedding
    import torch
    import torch.nn.functional as F
    empty_txt = ''
    clip_model = clip_vit_b32()
    empty_emb = get_embedding(clip_model, empty_txt)
    search_joined_values_emb = get_embedding(clip_model, search_joined_values)
    joined_values_emb = get_embedding(clip_model, joined_values)

    base_sim = F.cosine_similarity(
            torch.tensor(empty_emb).unsqueeze(0),
            torch.tensor(joined_values_emb).unsqueeze(0)
        ).item()
    search_sim = F.cosine_similarity(
            torch.tensor(search_joined_values_emb).unsqueeze(0),
            torch.tensor(joined_values_emb).unsqueeze(0)
        ).item()
    relative_sim = search_sim - base_sim

    print("base_sim(empty_txt_emb, joined_values_emb):", base_sim)
    print("search_sim(search_joined_values_emb, joined_values_emb):", search_sim)
    print("relative_sim (search_sim - base_sim):", relative_sim)

def test_matched_info_from_csv():
    import pandas as pd
    csv_file_path = '/mnt/nas/shengjie/qdrant_data/qwenvl3_data/qwenvl3_cloth_data.csv'
    df = pd.read_csv(csv_file_path)


    # ========== Step 1: 读取模板json字符串 ==========
    template_json_path = 'demo_qwenvl_clothing.json'
    with open(template_json_path, 'r', encoding='utf-8') as f:
        template_json_str = f.read()

    # ========== Step 2: 构建question + client + template = ans ==========
    search_txt = '春秋两季黑色大衣，两个口袋'

    qwen_client = get_qwen3_client()
    
    ans = get_ans_dict_by_search(qwen_client, template_json_str, search_txt)
    
    # ========== Step 3: ans => ans_dict ==========
    import json
    try:
        ans_dict = json.loads(ans)
    except Exception as e:
        print("Failed to parse ans as JSON:", e)
        ans_dict = None

    # ========== Step 4: ans_dict => non keys ==========
    non_empty_keys = get_non_empty_keys(ans_dict)

    # ========== Step 5: 构建 target_json_dict ==========

    # ========== Step 6: 获取 target_json_dict's concat value + ans_dict's concat value ==========
    # —— 从模板取内容并拼接
    search_joined_values_list = get_matched_info_by_jsondict(ans_dict,
                                                        non_empty_keys)
    search_joined_values = '&&'.join(search_joined_values_list)
    # print("从ans_dict收集的所有value并拼接:", search_joined_values)

    # —— 从模板取内容并拼接
    from util_for_clipvitb32 import clip_vit_b32, get_embedding
    import torch
    import torch.nn.functional as F
    import heapq
    clip_model = clip_vit_b32()
    empty_txt = ''
    empty_emb = get_embedding(clip_model, empty_txt)    
    search_joined_values_emb = get_embedding(clip_model, search_joined_values)

    topk = 10
    # Python的最小堆，存储(tuple: relative_sim, idx, img_path)
    heap = []

    for idx, row in df.iterrows():
        target_json_dict_str = row.get('desc_json', None)
        img_path = row.get('img_path', '')

        try:
            target_json_dict = json.loads(target_json_dict_str)
        except Exception as e:
            print("template_json_str 解析失败:", e)
            target_json_dict = {}

        joined_values_list = get_matched_info_by_jsondict(target_json_dict, non_empty_keys)
        joined_values = '&&'.join(joined_values_list)
        # print("收集的所有value并拼接:", joined_values)

        joined_values_emb = get_embedding(clip_model, joined_values)

        base_sim = F.cosine_similarity(
            torch.tensor(empty_emb).unsqueeze(0),
            torch.tensor(joined_values_emb).unsqueeze(0)
        ).item()
        search_sim = F.cosine_similarity(
            torch.tensor(search_joined_values_emb).unsqueeze(0),
            torch.tensor(joined_values_emb).unsqueeze(0)
        ).item()
        relative_sim = search_sim - base_sim

        # 只保留relative_sim>=0的数据
        if relative_sim >= 0:
            item = (relative_sim, search_sim, base_sim, idx, img_path, joined_values)
            if len(heap) < topk:
                heapq.heappush(heap, item)
            else:
                # 维护topk个最大relative_sim, 如果新来的比堆顶大则替换
                if relative_sim > heap[0][0]:
                    heapq.heappushpop(heap, item)

    # 按relative_sim降序展示topk
    print(f'search_joined_values_emb={search_joined_values}')
    heap_sorted = sorted(heap, key=lambda x: -x[0])
    for sim, search_sim, base_sim, idx, path, tar_value_str in heap_sorted:
        print(f"TOPK relative_sim={sim:.4f} search_sim={search_sim:.4f} base_sim={base_sim:.4f} \
                idx={idx} path= {path} desc={tar_value_str}")


if __name__ == '__main__':
    import argparse,os
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--cuda', type=str, default='2', help='CUDA device id')
    args, unknown = parser.parse_known_args()
    os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda
    # test()
    test_matched_info()
    # test_matched_info_from_csv()